"""
读取download_urls_only.json, 读取所有的download_url, 下载到本地
"""

import json
import os
import wget
from tqdm import tqdm
from urllib.parse import unquote
import concurrent.futures
import threading
from collections import deque

# 设置下载目录和已完成记录文件
DOWNLOAD_DIR = r"D:\Proj1126Files"
DONE_FILE = "done.json"
MAX_WORKERS = 5  # 同时下载的线程数
BATCH_SIZE = 10  # 累积多少个ID才写入文件
MAX_RETRIES = 1  # 最大重试次数
TIMEOUT = 10  # 请求超时时间
# 确保下载目录存在
os.makedirs(DOWNLOAD_DIR, exist_ok=True)

def sanitize_filename(filename):
    # 替换 Windows 文件系统中的非法字符
    illegal_chars = ['\\', '/', ':', '*', '?', '"', '<', '>', '|']
    for char in illegal_chars:
        filename = filename.replace(char, '_')
    return filename

# 读取已完成的下载记录
done_attach_ids = set()
if os.path.exists(DONE_FILE):
    try:
        with open(DONE_FILE, "r", encoding="utf-8") as f:
            done_attach_ids = set(json.load(f))
    except:
        print(f"读取{DONE_FILE}失败，将创建新文件")

# 读取待下载的URL
with open("download_urls_only.json", "r", encoding="utf-8") as f:
    download_urls = json.load(f)

# 添加线程锁和待写入队列
done_lock = threading.Lock()
tqdm_lock = threading.Lock()
pending_ids = deque(maxlen=BATCH_SIZE)

def write_pending_ids(force=False):
    """将待写入的ID写入文件"""
    global pending_ids
    with done_lock:
        if force or len(pending_ids) >= BATCH_SIZE:
            done_attach_ids.update(pending_ids)
            with open(DONE_FILE, "w", encoding="utf-8") as f:
                json.dump(list(done_attach_ids), f, ensure_ascii=False, indent=2)
            pending_ids.clear()

def download_file(item, pbar):
    """单个文件下载函数"""
    url = item["download_url"]
    attach_id = item["attach_id"]
    
    for retry in range(MAX_RETRIES):
        try:
            # 从URL中获取文件名
            if "?n=" in url:
                filename = unquote(url.split("?n=")[1])
            else:
                filename = url.split("/")[-1]
            filename = sanitize_filename(filename)
            
            # 构建保存路径
            save_path = os.path.join(DOWNLOAD_DIR, filename)
            print(f"下载文件: {filename} -> {save_path}")
            # 使用wget下载文件
            wget.download(url, save_path)
            
            # 添加到待写入队列
            with done_lock:
                pending_ids.append(attach_id)
                write_pending_ids()
            
            # 更新进度条
            with tqdm_lock:
                pbar.update(1)
            
            # 下载成功，跳出重试循环
            break
                    
        except Exception as e:
            if retry < MAX_RETRIES - 1:
                print(f"\n下载失败: {url}")
                print(f"错误信息: {str(e)}")
                print(f"正在进行第{retry + 2}次尝试...")
                # 如果文件已经部分下载，删除它
                if os.path.exists(save_path):
                    try:
                        os.remove(save_path)
                    except:
                        pass
            else:
                print(f"\n下载失败: {url}")
                print(f"错误信息: {str(e)}")
                print(f"已达到最大重试次数，跳过该文件")
                with tqdm_lock:
                    pbar.update(1)

def main():
    # 过滤掉已下载的文件
    remaining_downloads = [item for item in download_urls if item["attach_id"] not in done_attach_ids]
    print(f"总计{len(download_urls)}个文件，已完成{len(done_attach_ids)}个，剩余{len(remaining_downloads)}个待下载")

    # 创建进度条
    with tqdm(total=len(remaining_downloads), desc="下载文件") as pbar:
        # 使用线程池执行下载任务
        with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
            # 提交所有下载任务
            futures = [
                executor.submit(download_file, item, pbar)
                for item in remaining_downloads
            ]
            # 等待所有任务完成，如果超时，则强制写入剩余的pending_ids
            concurrent.futures.wait(futures, timeout=TIMEOUT)
            write_pending_ids(force=True)

if __name__ == "__main__":
    main()